from keras.datasets import cifar100
from keras import Sequential, layers, activations, optimizers, losses, Model

(x_train, y_train), (x_test, y_test) = cifar100.load_data()
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0


class VGG16(Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv = Sequential([
            layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same'),
            layers.Conv2D(filters=16, kernel_size=(3, 3), padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same'),
            layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same'),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same'),
            layers.Conv2D(filters=64, kernel_size=(3, 3), padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same'),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same'),
            layers.Conv2D(filters=128, kernel_size=(3, 3), padding='same'),
            layers.MaxPooling2D()
        ])
        self.flat = Sequential([layers.Flatten()])
        self.fc = Sequential([
            layers.Dense(units=512, activation=activations.relu),
            layers.Dropout(0.4),
            layers.Dense(units=512, activation=activations.relu),
            layers.Dropout(0.4),
            layers.Dense(units=100, activation=activations.softmax)
        ])

    def call(self, inputs, training=None, mask=None):
        out = self.conv(inputs)
        out = self.flat(out)
        out = self.fc(out)
        return out


model = VGG16()
model.build(input_shape=(None, 32, 32, 3))
model.summary()
model.compile(optimizer=optimizers.Adam(), loss=losses.sparse_categorical_crossentropy, metrics='acc')
model.fit(x_train, y_train, batch_size=100, epochs=2)

model.evaluate(x_test, y_test)
